{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NpJd3dlOCStH"
      },
      "source": [
        "\u003ca href=\"https://colab.research.google.com/github/magenta/ddsp/blob/master/ddsp/colab/tutorials/3_training.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hMqWDc_m6rUC"
      },
      "source": [
        "\n",
        "##### Copyright 2020 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "VNhgka4UKNjf"
      },
      "outputs": [],
      "source": [
        "# Copyright 2020 Google LLC. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ZFIqwYGbZ-df"
      },
      "source": [
        "# DDSP Training\n",
        "\n",
        "This notebook demonstrates the libraries in [https://github.com/magenta/ddsp/tree/master/ddsp/training](https://github.com/magenta/ddsp/tree/master/ddsp/training). It is a simple example, overfitting a single audio sample, for educational purposes. \n",
        "\n",
        "_For a full training pipeline please use [ddsp/training/ddsp_run.py](https://github.com/magenta/ddsp/blob/master/ddsp/training/README.md#train-1) as in the [train_autoencoder.ipynb](https://github.com/magenta/ddsp/blob/master/ddsp/colab/demos/train_autoencoder.ipynb)_.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "S_jXCnwZ2QYW"
      },
      "outputs": [],
      "source": [
        "# Install and import dependencies\n",
        "%tensorflow_version 2.x\n",
        "!pip install -qU ddsp\n",
        "\n",
        "# Ignore a bunch of deprecation warnings\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "\n",
        "import time\n",
        "\n",
        "import ddsp\n",
        "from ddsp.training import (data, decoders, encoders, models, preprocessing, \n",
        "                           train_util, trainers)\n",
        "from ddsp.colab.colab_utils import play, specplot, DEFAULT_SAMPLE_RATE\n",
        "import gin\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import tensorflow.compat.v2 as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "sample_rate = DEFAULT_SAMPLE_RATE  # 16000"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "khYj8yiMDxGL"
      },
      "source": [
        "# Get a batch of data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "IzzaWKxVkYms"
      },
      "outputs": [],
      "source": [
        "# Get a single example from NSynth.\n",
        "# Takes a few seconds to load from GCS.\n",
        "data_provider = data.NSynthTfds(split='test')\n",
        "dataset = data_provider.get_batch(batch_size=1, shuffle=False).take(1).repeat()\n",
        "batch = next(iter(dataset))\n",
        "audio = batch['audio']\n",
        "n_samples = audio.shape[1]\n",
        "\n",
        "specplot(audio)\n",
        "play(audio)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "acGO0ifg0I3k"
      },
      "source": [
        "# Get a distribution strategy\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "VUlAeW40Wsvr"
      },
      "outputs": [],
      "source": [
        "strategy = train_util.get_strategy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Op0V8onI0VUK"
      },
      "source": [
        "# Get model and trainer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "EWZQXFLehCU0"
      },
      "source": [
        "## python "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "HCqXRY1KeX8S"
      },
      "outputs": [],
      "source": [
        "TIME_STEPS = 1000\n",
        "\n",
        "# Create Neural Networks.\n",
        "preprocessor = preprocessing.DefaultPreprocessor(time_steps=TIME_STEPS)\n",
        "\n",
        "decoder = decoders.RnnFcDecoder(rnn_channels = 256,\n",
        "                                rnn_type = 'gru',\n",
        "                                ch = 256,\n",
        "                                layers_per_stack = 1,\n",
        "                                input_keys = ('ld_scaled', 'f0_scaled'),\n",
        "                                output_splits = (('amps', 1),\n",
        "                                                 ('harmonic_distribution', 45),\n",
        "                                                 ('noise_magnitudes', 45)))\n",
        "\n",
        "# Create Processors.\n",
        "additive = ddsp.synths.Additive(n_samples=n_samples, \n",
        "                                sample_rate=sample_rate,\n",
        "                                name='additive')\n",
        "\n",
        "noise = ddsp.synths.FilteredNoise(window_size=0,\n",
        "                                  initial_bias=-10.0,\n",
        "                                  name='noise')\n",
        "add = ddsp.processors.Add(name='add')\n",
        "\n",
        "# Create ProcessorGroup.\n",
        "dag = [(additive, ['amps', 'harmonic_distribution', 'f0_hz']),\n",
        "       (noise, ['noise_magnitudes']),\n",
        "       (add, ['noise/signal', 'additive/signal'])]\n",
        "\n",
        "processor_group = ddsp.processors.ProcessorGroup(dag=dag,\n",
        "                                                 name='processor_group')\n",
        "\n",
        "\n",
        "# Loss_functions\n",
        "spectral_loss = ddsp.losses.SpectralLoss(loss_type='L1',\n",
        "                                         mag_weight=1.0,\n",
        "                                         logmag_weight=1.0)\n",
        "\n",
        "with strategy.scope():\n",
        "  # Put it together in a model.\n",
        "  model = models.Autoencoder(preprocessor=preprocessor,\n",
        "                             encoder=None,\n",
        "                             decoder=decoder,\n",
        "                             processor_group=processor_group,\n",
        "                             losses=[spectral_loss])\n",
        "  trainer = trainers.Trainer(model, strategy, learning_rate=1e-3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "uAZgDMV9hGyp"
      },
      "source": [
        "## or [`gin`](https://github.com/google/gin-config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1JPmTwQshVya"
      },
      "outputs": [],
      "source": [
        "gin_string = \"\"\"\n",
        "import ddsp\n",
        "import ddsp.training\n",
        "\n",
        "# Preprocessor\n",
        "models.Autoencoder.preprocessor = @preprocessing.DefaultPreprocessor()\n",
        "preprocessing.DefaultPreprocessor.time_steps = 1000\n",
        "\n",
        "\n",
        "# Encoder\n",
        "models.Autoencoder.encoder = None\n",
        "\n",
        "# Decoder\n",
        "models.Autoencoder.decoder = @decoders.RnnFcDecoder()\n",
        "decoders.RnnFcDecoder.rnn_channels = 256\n",
        "decoders.RnnFcDecoder.rnn_type = 'gru'\n",
        "decoders.RnnFcDecoder.ch = 256\n",
        "decoders.RnnFcDecoder.layers_per_stack = 1\n",
        "decoders.RnnFcDecoder.input_keys = ('ld_scaled', 'f0_scaled')\n",
        "decoders.RnnFcDecoder.output_splits = (('amps', 1),\n",
        "                                       ('harmonic_distribution', 20),\n",
        "                                       ('noise_magnitudes', 20))\n",
        "\n",
        "# ProcessorGroup\n",
        "models.Autoencoder.processor_group = @processors.ProcessorGroup()\n",
        "\n",
        "processors.ProcessorGroup.dag = [\n",
        "  (@additive/synths.Additive(),\n",
        "    ['amps', 'harmonic_distribution', 'f0_hz']),\n",
        "  (@noise/synths.FilteredNoise(),\n",
        "    ['noise_magnitudes']),\n",
        "  (@add/processors.Add(),\n",
        "    ['noise/signal', 'additive/signal']),\n",
        "]\n",
        "\n",
        "# Additive Synthesizer\n",
        "additive/synths.Additive.name = 'additive'\n",
        "additive/synths.Additive.n_samples = 64000\n",
        "additive/synths.Additive.scale_fn = @core.exp_sigmoid\n",
        "\n",
        "# Filtered Noise Synthesizer\n",
        "noise/synths.FilteredNoise.name = 'noise'\n",
        "noise/synths.FilteredNoise.n_samples = 64000\n",
        "noise/synths.FilteredNoise.window_size = 0\n",
        "noise/synths.FilteredNoise.scale_fn = @core.exp_sigmoid\n",
        "noise/synths.FilteredNoise.initial_bias = -10.0\n",
        "\n",
        "# Add\n",
        "add/processors.Add.name = 'add'\n",
        "\n",
        "models.Autoencoder.losses = [\n",
        "    @losses.SpectralLoss(),\n",
        "]\n",
        "losses.SpectralLoss.loss_type = 'L1'\n",
        "losses.SpectralLoss.mag_weight = 1.0\n",
        "losses.SpectralLoss.logmag_weight = 1.0\n",
        "\"\"\"\n",
        "\n",
        "with gin.unlock_config():\n",
        "  gin.parse_config(gin_string)\n",
        "\n",
        "with strategy.scope():\n",
        "  # Autoencoder arguments are filled by gin.\n",
        "  model = ddsp.training.models.Autoencoder()\n",
        "  trainer = trainers.Trainer(model, strategy, learning_rate=1e-3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MnnxpYbRrPrp"
      },
      "source": [
        "# Train"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "IGvCE5BbrWTU"
      },
      "source": [
        "## Build model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1lQW604_QWm1"
      },
      "outputs": [],
      "source": [
        "# Build model, easiest to just run forward pass.\n",
        "dataset = trainer.distribute_dataset(dataset)\n",
        "trainer.build(next(iter(dataset)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "RFEqt6e1DsqG"
      },
      "source": [
        "## Train Loop"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "LWdoRIONDxri"
      },
      "outputs": [],
      "source": [
        "dataset_iter = iter(dataset)\n",
        "\n",
        "for i in range(300):\n",
        "  losses = trainer.train_step(dataset_iter)\n",
        "  res_str = 'step: {}\\t'.format(i)\n",
        "  for k, v in losses.items():\n",
        "    res_str += '{}: {:.2f}\\t'.format(k, v)\n",
        "  print(res_str)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "2cj220vSF8_Y"
      },
      "source": [
        "# Analyze results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mDU_FysURm_Z"
      },
      "outputs": [],
      "source": [
        "# Run a batch of predictions.\n",
        "start_time = time.time()\n",
        "controls =  model(next(dataset_iter))\n",
        "audio_gen = model.get_audio_from_outputs(controls)\n",
        "print('Prediction took %.1f seconds' % (time.time() - start_time))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "AlE5kkroHiFr"
      },
      "outputs": [],
      "source": [
        "print('Original Audio')\n",
        "play(audio)\n",
        "print('Resynthesized Audio')\n",
        "play(audio_gen)\n",
        "print('Filtered Noise Audio')\n",
        "audio_noise = controls['noise']['signal']\n",
        "play(audio_noise)\n",
        "\n",
        "specplot(audio)\n",
        "specplot(audio_gen)\n",
        "specplot(audio_noise)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "DVhoLzV-ZYav"
      },
      "outputs": [],
      "source": [
        "batch_idx = 0\n",
        "get = lambda key: ddsp.core.nested_lookup(key, controls)[batch_idx]\n",
        "\n",
        "amps = get('additive/controls/amplitudes')\n",
        "harmonic_distribution = get('additive/controls/harmonic_distribution')\n",
        "noise_magnitudes = get('noise/controls/magnitudes')\n",
        "f0_hz = get('f0_hz')\n",
        "loudness = get('loudness_db')\n",
        "\n",
        "audio_noise = get('noise/signal')\n",
        "\n",
        "f, ax = plt.subplots(1, 2, figsize=(14, 4))\n",
        "f.suptitle('Input Features', fontsize=16)\n",
        "ax[0].plot(loudness)\n",
        "ax[0].set_ylabel('Loudness')\n",
        "ax[1].plot(f0_hz)\n",
        "ax[1].set_ylabel('F0_Hz')\n",
        "\n",
        "f, ax = plt.subplots(1, 2, figsize=(14, 4))\n",
        "f.suptitle('Synth Params', fontsize=16)\n",
        "ax[0].semilogy(amps)\n",
        "ax[0].set_ylabel('Amps')\n",
        "ax[0].set_ylim(1e-5, 2)\n",
        "# ax[0].semilogy(harmonic_distribution)\n",
        "ax[1].matshow(np.rot90(np.log10(harmonic_distribution + 1e-6)),\n",
        "              cmap=plt.cm.magma, \n",
        "              aspect='auto')\n",
        "ax[1].set_ylabel('Harmonic Distribution')\n",
        "ax[1].set_xticks([])\n",
        "_ = ax[1].set_yticks([])\n",
        "\n",
        "f, ax = plt.subplots(1, 1, figsize=(7, 4))\n",
        "# f.suptitle('Filtered Noise Params', fontsize=16)\n",
        "ax.matshow(np.rot90(np.log10(noise_magnitudes + 1e-6)), \n",
        "           cmap=plt.cm.magma, \n",
        "           aspect='auto')\n",
        "ax.set_ylabel('Filtered Noise Magnitudes')\n",
        "ax.set_xticks([])\n",
        "_ = ax.set_yticks([])\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "hMqWDc_m6rUC",
        "ZFIqwYGbZ-df",
        "khYj8yiMDxGL",
        "acGO0ifg0I3k",
        "Op0V8onI0VUK",
        "EWZQXFLehCU0",
        "uAZgDMV9hGyp",
        "MnnxpYbRrPrp",
        "IGvCE5BbrWTU",
        "RFEqt6e1DsqG",
        "2cj220vSF8_Y"
      ],
      "last_runtime": {},
      "name": "3_training.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
